//============================================================================
// Name        : Coevolution 4
// Author      : Alex Seeholzer
// Description :
// Downgrade of Coevolution 3. Removing beta fitness level, and treating all is particles except the ones
// driven by the virus equally, i.e. assigning base fitness to all.
// We initialize 2 fixed peaks in virus fitness, seed only for random # generation.
//============================================================================
// cin with strings

#include <gsl/gsl_math.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_interp.h>


#define TYPE unsigned char
#define L 8
#define ONE 1
#define RRMASK 0x000000FF
#define MAXINT 0xFFFFFFFF
#define LRR 8
#define VIR 0
#define IS 1
#define DELTA 1.
#define N 300000
#define TMAX 30000

#define NM 2 /* population number, 0=virus 1=immune */
bool vir_masters[(ONE<<LRR)]; // array of master sequences vor virus functional region, true == master

double alpha;
double delta;
double gam[2];
double initratio;
int tmax;

TYPE popnew[NM][N];
double sel[N];
double mut[N];
double selectr[2*N];

inline TYPE mutate (TYPE x, double mu);
inline int hd (TYPE x, TYPE y); // hamming dist
inline void outbinary(TYPE num);
void reproduce (TYPE pop[NM][N], double* mu, unsigned int cdis[NM][(ONE << LRR)]);
void sample(TYPE* pop,TYPE* popnew,double* mut,double* sel,double mu);
TYPE compute_consensus (TYPE *pop);

// computes for a pop the distribution of clones (stored in cdis)
void compute_cdis (TYPE *pop, unsigned int *cdis);

// Random number generator
gsl_rng *r;
gsl_interp_accel *acc;
const gsl_rng_type *T;

int main(int argc, char **argv) {

	// Set up output file
	FILE * outfile;
	int i,j,k,t,s;
	TYPE pop[NM][N]; // array of population
	double mu[NM]; // mutation rate
	TYPE smax[NM]; // maximal populated sequences for each population
	unsigned int cdis[NM][(ONE << LRR)]; // for each population, distribution of functional region

	// random number generators
	gsl_rng_env_setup();
	acc = gsl_interp_accel_alloc();
	T = gsl_rng_gfsr4;
	r = gsl_rng_alloc(T);
	bool redoutput=false;

	// read mu, s = seed
	if (argc > 9) {
		mu[VIR]=atof(argv[1]);
		mu[IS]=atof(argv[2]);
		gam[VIR]=atof(argv[3]);
		gam[IS]=atof(argv[4]);
		delta=atof(argv[5]);
		alpha=atof(argv[6]);
		initratio=atof(argv[7]);
		tmax=atoi(argv[8]);
		s=atoi(argv[9]);
		if (argc > 10)
			redoutput=atoi(argv[10]);
		fprintf(stderr, "\nmu_v=% .3e, mu_is=% .3e, gamma_v=% .3e, gamma_is=% .3e, delta=% .3e, alpha=% .3e, initratio=%.3e, tmax=%i, s=%i, N= %i, reduced output=%i\n", mu[VIR], mu[IS], gam[VIR], gam[IS], delta, alpha, initratio, tmax, s, N, redoutput);
	} else {
		fprintf(stderr, "not enough args. argc=%i\n", argc);
		exit(1);
	}

	char *outfilename = (char*) malloc(255);
	sprintf(outfilename, "runMod4_%s_%s_%s_%s_%s_%s_%s_%s_%s_%i.dat", argv[1],argv[2],argv[3],argv[4],argv[5],argv[6],argv[7],argv[8],argv[9],N);
	outfile = fopen(outfilename,"w");
	free(outfilename);
	fprintf(outfile, "mu_v=% .3e, mu_is=% .3e, gamma_v=% .3e, gamma_is=% .3e, delta=% .3e, alpha=% .3e, initratio=%.3e, tmax=%i, s=%i, N= %i, reduced output=%i\n", mu[VIR], mu[IS], gam[VIR], gam[IS], delta, alpha, initratio, tmax, s, N,redoutput);
	gsl_rng_set (r, s); // initialize random num generator with seed s

	// compute vir_masters_num unique fitness peaks in [0,(ONE<<LRR)] and store in vir_masters
	const int vir_masters_num = 2;
	int vir_masters_idx[vir_masters_num]={231,255};
	for(i=0;i<(ONE<<LRR);i++)
		vir_masters[i]=false;
	for(i=0;i<vir_masters_num;i++)
		vir_masters[vir_masters_idx[i]]=true;
	unsigned int ham = hd(vir_masters_idx[0],vir_masters_idx[1]);
	fprintf(stderr,"\nMaster sequences for virus population: (hamdist = %i)\n",ham);
	fprintf(outfile,"Master sequences for virus population: (hamdist = %i)\n",ham);
	for(i=0;i<vir_masters_num;i++) {
		fprintf(stderr,"%i ",vir_masters_idx[i]);
		fprintf(outfile,"%i ",vir_masters_idx[i]);
	}
	fprintf(outfile,"\n");
	fprintf(stderr,"\n");
	fprintf(stderr,"\n");

	// immune population initialization, set as initratio * first, (1-initratio) * second
	for (i = 0; i < N; i++) {
		int mstr = (i>=N*initratio) ?  vir_masters_idx[0] : vir_masters_idx[1];
		pop[IS][i] = mstr;
	}

	// virus population initialization, set as initratio * first, (1-initratio) * second
	for (i = 0; i < N; i++) {
		int mstr = (i>=N*initratio) ?  vir_masters_idx[0] : vir_masters_idx[1];
		pop[VIR][i] = mstr;
	}

	// compute frequency of master sequences for initial reproduce call
	for (j = 0; j < NM; j++)
		compute_cdis(pop[j], cdis[j]);

	// main loop - reproduce, calc distribution and output
	for (t = 0; t < tmax; t++) {

		reproduce(pop, mu, cdis);

		for (j = 0; j < NM; j++)
			compute_cdis(pop[j], cdis[j]);
		if(!redoutput)
			for (k = 0; k < (ONE << LRR); k++) {
				fprintf(outfile,"%i %i ", t, k);
				for (j = 0; j < NM; j++)
					fprintf(outfile,"%i ", cdis[j][k]);
				fprintf(outfile,"\n");
			}
		else
			for (k = 0; k < vir_masters_num; k++) {
				int gene = vir_masters_idx[k];
				fprintf(outfile,"%i %i ", t, gene);
				for (j = 0; j < NM; j++)
					fprintf(outfile,"%i ", cdis[j][gene]);
				fprintf(outfile,"\n");
			}

		if (t%100 == 0)
			fprintf(stderr,"\r%.0f%%",t/(double)tmax*100.);
	}

	fprintf(stderr,"\n");
	fprintf(stderr,"\n");

	for (j = 0; j < NM; j++) {
		fprintf(stderr,"Population %i\n",j);
		smax[j] = 0;
		for (k = 1; k < (ONE << LRR); k++)
			if (cdis[j][k] > cdis[j][smax[j]])
				smax[j] = k;
		fprintf(stderr, "smax[%i]= %3i (%i), neighbors: ", j, smax[j],
				cdis[j][smax[j]]);
		for (k = 0; k < LRR; k++)
			fprintf(stderr, "%3i (%i), ", smax[j] ^ (ONE << k), cdis[j][(smax[j] ^ (ONE << k))]);
		fprintf(stderr, "\n");
	}

	fclose(outfile);
	return 0;
}

void reproduce(TYPE pop[NM][N], double* mu, unsigned int cdis[NM][(ONE << LRR)]) {
	int i, j;
	double D;

	// reproduce virus population
	for (i = 0; i < N; i++) {
		// calculate percentage of genome pop[VIR][i] in pop[IS]
		D = cdis[IS][(unsigned int)pop[VIR][i]] / (double) N;
		mut[i] = delta;
		sel[i] = (1/DELTA)-(gam[VIR]*D);
		if(sel[i]<0) {
			fprintf(stderr,"<0");
			sel[i]=0;
		}
		if (vir_masters[pop[VIR][i]]) // if we have one of the master sequences
			mut[i] = alpha;
	}

	sample(pop[VIR],popnew[VIR],mut,sel,mu[VIR]);

	// reproduce immune system population
	for (i = 0; i < N; i++) {
		// calculate percentage of genome pop[IS][i] in pop[VIR]
		D = cdis[VIR][(unsigned int)pop[IS][i]] / (double) N;
		sel[i] = (1/DELTA);
		mut[i] = delta+(gam[IS]*D);
	}
	sample(pop[IS],popnew[IS],mut,sel,mu[IS]);

	for (j = 0; j < NM; j++)
		for (i = 0; i < N; i++)
			pop[j][i] = popnew[j][i];
	return;
}
void sample(TYPE* pop, TYPE* popnew, double* mut,double* sel,double mu) {
	double ptot;
	int i,index;
	// sum of probabilities for normalization
	ptot = 0;

	for (i = 0; i < N; i++)
		ptot += sel[i] + mut[i];

	// evaluate probabilities to be selected
	// (even index) and to be selected+mutated
	// (odd indices)
	selectr[0] = sel[0] / ptot;
	selectr[1] = mut[0] / ptot + selectr[0];
	for (i = 1; i < N; i++) {
		selectr[2* i ] = sel[i] / ptot + selectr[2* i - 1];
		selectr[2* i + 1] = mut[i] / ptot + selectr[2* i ];
	}
	//get N random numbers in [0,1)
	for (i = 0;	i < N;i++) {
		index = gsl_interp_bsearch(selectr, gsl_rng_uniform(r), 0, 2*N - 2) + 1; //look up entry in select array

		// if index is odd, mutate this sequence with probability mu, otherwise copy
		popnew[i] = (index % 2 == 1 ? mutate(pop[index / 2], mu) : pop[index / 2]);
	}
	return;
}

inline TYPE mutate (TYPE x, double mu) {
  int i;
  TYPE out=x;
  for (i=0; i<L; i++)
    if (gsl_rng_uniform(r) < mu)
      out^=(ONE << i);
  return out;
}

inline void outbinary(TYPE num) {
	TYPE tmp = num;
	for(int i=0;i<L;i++) {
		int powr = int(pow(2,(L-i-1)));
		int res = tmp/powr;
		fprintf(stderr,"%i",res);
		tmp-=res*powr;
	}
	fprintf(stderr,"\n");
}

inline int hd (TYPE x, TYPE y) {
  int dist=0;
  TYPE val=x ^ y;
  while (val) {
    ++dist;
    val &= val-1;
  }
  return dist;
}

void compute_cdis (TYPE *pop, unsigned int *cdis) {
  for (int k=0; k<(ONE << LRR); k++)
    cdis[k]=0;
  for (int i=0; i<N; i++) {
	  cdis[(unsigned int)pop[i]]++;
  }
  return;
}
